import os

from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langgraph.prebuilt import create_react_agent, ToolNode

base_url = os.environ.get("DASHSCOPE_BASE_URL")
model_name = os.environ.get("DASHSCOPE_MODEL_NAME")
os.environ["OPENAI_API_KEY"] = os.getenv("DASHSCOPE_API_KEY")

llm = ChatOpenAI(base_url=base_url, model_name=model_name)


@tool("devide_tool")
def devide(a: int, b: int) -> float:
    """
    计算两个整数的除法。
    Args：
        a（int）：除数
        b（int）：被除数
    """
    print("call devide")
    print("a: ", a)
    print("b: ", b)
    if b == 1:
        raise ValueError("除数不能为1")
    if b == 0:
        raise ZeroDivisionError("除数不能为0")
    return a / b


# 定义工具调用错误时的处理函数
def handle_tool_error(error: Exception) -> str:
    """
    处理工具调用错误
    Args：
        error（Exception）：工具调用错误
    """
    print("handle tool error")
    if isinstance(error, ValueError):
        return "除数为1没有意义，请重新输入一个除数和被除数。"
    elif isinstance(error, ZeroDivisionError):
        return "除数不能为0，请重新输入一个除数和被除数。"
    return f"工具调用错误：{error}"


tool_node = ToolNode(
    [devide],
    handle_tool_errors=handle_tool_error
)

agent_with_tool_error_handler = create_react_agent(
    model=llm,
    tools=tool_node,
    prompt="请根据tool回答问题，如果调用tool发生错误，则按照错误处理函数给出错误提示信息。请让tool来判断应该如何处理，而不是自己处理并绕过tool直接返回。"
)

result = agent_with_tool_error_handler.invoke(
    {"messages": [{"role": "user", "content": "10除以0等于多少？"}]}
)
print(result)
